import os
from openai import OpenAI

client = OpenAI(
    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
    api_key="sk-30a77ee37da8436596edbc57242afc25",
    # 如何获取API Key：https://help.aliyun.com/zh/model-studio/developer-reference/get-api-key
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

question = """
你是谁
"""

stream_response = client.chat.completions.create(
    model="deepseek-r1",
    # model="deepseek-v3",
    messages=[
        {'role': 'system', 'content': 'You are a helpful assistant.'},
        {'role': 'user', 'content': question}
    ],
    stream=True,
    stream_options={
        "include_usage": True
    }
)

output_text = ""
reasoning_content = ""
content = ""
completion_tokens = 0
prompt_tokens = 0
for chunk in stream_response:
    print(chunk)
    # 判断思考包
    if len(chunk.choices) > 0 and chunk.choices[0].delta.reasoning_content is not None:
        reasoning_content += chunk.choices[0].delta.reasoning_content
        print(chunk.choices[0].delta.reasoning_content)

    elif len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None:
        # output_text += chunk.choices[0].delta.content
        # 打印输出内容
        print(chunk.choices[0].delta.content, end="")
        content += chunk.choices[0].delta.content
    else:
        completion_tokens = chunk.usage.completion_tokens
        prompt_tokens = chunk.usage.prompt_tokens
    # print(output_text)

# # print(stream_response.choices[0].message.content)
# print("输入 Tokens:", stream_response.usage.prompt_tokens)
# print("输出 Tokens:", stream_response.usage.completion_tokens)
# print("总 Tokens:", stream_response.usage.total_tokens)

print(reasoning_content)
print(content)
print(prompt_tokens)
print(completion_tokens)
